#include <stdio.h>
#include <math.h>
#include <vector>
#include <stack>
#include<iostream>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

#include "lane_decode.h"

namespace py = pybind11;

int x_mv[] = { 0, 1, 0, -1 };
int y_mv[] = { -1, 0, 1, 0 };

inline float sigmoid(float const x)  { return (1 / (1 + exp(-x))); };

inline bool same_edge(int const* direction_0,
                      int const* direction_1)
{
    return std::fabs(direction_0[0] +  direction_1[0]) < 1e-6 && std::fabs(direction_0[1] +  direction_1[1]) < 1e-6;
};

inline bool same_point(float const* xy_0,
                      float const* xy_1)
{
    return std::fabs(xy_0[0] - xy_1[0]) < 0.3 || std::fabs(xy_0[1] - xy_1[1]) < 0.3;
};

inline void distance_decode(float const distance,
                                    bool const is_cw,
                                    int* direction,
                                    float* xy)
{
    float cw_distance{distance};
    if (!is_cw)
    {
        cw_distance = 4 - cw_distance;
    }

    if (cw_distance >= 0 && cw_distance <= 1)
    {
        direction[0] =  0;
        direction[1] =  -1;
        xy[0] = cw_distance;
        xy[1] = 0;
    }
    else if (cw_distance > 1 && cw_distance <= 2)
    {
        direction[0] =  1;
        direction[1] =  0;
        xy[0] = 1;
        xy[1] = cw_distance - 1;
    }
    else if (cw_distance > 2 && cw_distance <= 3)
    {
        direction[0] =  0;
        direction[1] =  1;
        xy[0] = 3 - cw_distance;
        xy[1] = 1;
    }
    else if (cw_distance > 3 && cw_distance <= 4)
    {
        direction[0] =  -1;
        direction[1] =  0;
        xy[0] = 0;
        xy[1] = 4 - cw_distance;
    }
    else
    {
        direction[0] =  0;
        direction[1] =  -1;
        xy[0] = 0;
        xy[1] = 0;
    }
}


int dfs(std::uint16_t const start_x,
        std::uint16_t const start_y,
        std::uint16_t const start_pos,
        std::uint16_t const start_cls,
        int const& id,
        float const* dpuout_cls,
        float const* dpuout_reg,
        float const* class_threshold,
        int const& max_result_size,
        int filled_size_,
        int const* type_map,
        int const& net_length,
        int const& net_width,
        int const& preds_per_voxel,
        int const& cls_number,
        int const& reg_number,
        int const& lane_feature_size,
        float const* voxels_size,
        float const* pts_range,
        bool* visited,
        float* dense_pts_ptr)
{   
    int reg_pos{type_map[start_cls]};
    std::stack<std::tuple<std::uint16_t, std::uint16_t, std::uint16_t>> grids;
    grids.push(std::make_tuple(start_x, start_y, start_pos));

    std::uint64_t cur_index, next_index;
    std::uint16_t cur_x, cur_y, cur_pos;
    int cur_direction[2]{0};
    float cur_xy[2]{0.0};
    int next_direction[2]{0};
    float next_xy[2]{0.0};
    while (!grids.empty())
    {
        cur_x = std::get<0>(grids.top());
        cur_y = std::get<1>(grids.top());
        cur_pos = std::get<2>(grids.top());
        grids.pop();

        if (filled_size_ >= max_result_size)
        {
          return filled_size_;
        }

        cur_index =  (cur_pos * cls_number + start_cls) * net_length * net_width  + cur_x * net_width + cur_y;
        const float& cur_score = dpuout_cls[cur_index];
        visited[cur_index] = true;

        for (std::uint16_t i{0}; i < 2; ++i)
        {
            bool flag{true};
            const float& cur_distance = dpuout_reg[(cur_pos * reg_number * 2 + reg_pos * 2 + i) * net_length * net_width  + cur_x * net_width + cur_y];
            distance_decode(cur_distance, 1-i, cur_direction, cur_xy);
            for (std::uint16_t j{0}; j < 2; ++j)
            {
                next_index = (j * cls_number + start_cls) * net_length * net_width  + (cur_x + cur_direction[0]) * net_width + (cur_y + cur_direction[1]);
                const float& next_score = dpuout_cls[next_index];
                if (visited[next_index])
                {
                    flag = false;
                    continue;
                }
//                if (next_score < class_threshold[start_cls])
//                {
//                    continue;
//                }

//                std::cout << "cur_x: " << cur_x << std::endl;
//                std::cout << "cur_y: " << cur_y << std::endl;
//                std::cout << "cur_pos: " << cur_pos << std::endl;
//                std::cout << "reg_pos: " << reg_pos << std::endl;
//                std::cout << "cw: " << 1-i << std::endl;
//                std::cout << "cur_distance: " << cur_distance << std::endl;
//                std::cout << "cur_direction[0]: " << cur_direction[0] << std::endl;
//                std::cout << "cur_direction[1]: " << cur_direction[1] << std::endl;
//                std::cout << "cur_xy[0]: " << cur_xy[0] << std::endl;
//                std::cout << "cur_xy[1]: " << cur_xy[1] << std::endl;


                for (std::uint16_t k{0}; k < 2; ++k)
                {
                    const float& next_distance = dpuout_reg[(j * reg_number * 2 + reg_pos * 2 + k) * net_length * net_width  + (cur_x + cur_direction[0]) * net_width + (cur_y + cur_direction[1])];
                    distance_decode(next_distance, 1-k, next_direction, next_xy);



//                    std::cout << "cur_x + cur_direction[0]: " << cur_x + cur_direction[0] << std::endl;
//                    std::cout << "cur_y + cur_direction[1]: " << cur_y + cur_direction[1] << std::endl;
//                    std::cout << "next_pos: " << j << std::endl;
//                    std::cout << "reg_pos: " << reg_pos << std::endl;
//                    std::cout << "cw: " << 1 - k << std::endl;
//                    std::cout << "next_distance: " << next_distance << std::endl;
//                    std::cout << "next_direction[0]: " << next_direction[0] << std::endl;
//                    std::cout << "next_direction[1]: " << next_direction[1] << std::endl;
//                    std::cout << "next_xy[0]: " << next_xy[0] << std::endl;
//                    std::cout << "next_xy[1]: " << next_xy[1] << std::endl;
//
//
//                    std::cout << "same_edge(cur_direction, next_direction): " << same_edge(cur_direction, next_direction) << std::endl;
//                    std::cout << "same_point(cur_xy, next_xy): " << same_point(cur_xy, next_xy) << std::endl;

                    if (same_edge(cur_direction, next_direction) && same_point(cur_xy, next_xy)) //同一条边且距离够近
                    {
                        flag = false;

                        dense_pts_ptr[filled_size_ * lane_feature_size + 0] = static_cast<float>(id);
                        dense_pts_ptr[filled_size_ * lane_feature_size + 1] = (cur_y + cur_xy[1]) * voxels_size[0] + pts_range[0];
                        dense_pts_ptr[filled_size_ * lane_feature_size + 2] = (cur_x + cur_xy[0]) * voxels_size[1] + pts_range[1];
                        dense_pts_ptr[filled_size_ * lane_feature_size + 3] = 0;
                        dense_pts_ptr[filled_size_ * lane_feature_size + 4] = static_cast<float>(start_cls);
                        dense_pts_ptr[filled_size_ * lane_feature_size + 5] = sigmoid(cur_score);
                        filled_size_++;

                        grids.push(std::make_tuple(cur_x + cur_direction[0], cur_y + cur_direction[1], j));

                        break;
                    }
                }

                if (!flag)
                {
                    break;
                }
            }

            if (flag)
            {
                dense_pts_ptr[filled_size_ * lane_feature_size + 0] = static_cast<float>(id);
                dense_pts_ptr[filled_size_ * lane_feature_size + 1] = (cur_y + cur_xy[1]) * voxels_size[0] + pts_range[0];
                dense_pts_ptr[filled_size_ * lane_feature_size + 2] = (cur_x + cur_xy[0]) * voxels_size[1] + pts_range[1];
                dense_pts_ptr[filled_size_ * lane_feature_size + 3] = 0;
                dense_pts_ptr[filled_size_ * lane_feature_size + 4] = static_cast<float>(start_cls);
                dense_pts_ptr[filled_size_ * lane_feature_size + 5] = sigmoid(cur_score);
                filled_size_++;
            }
        }
    }
    return filled_size_;
}

int decode_lane_double_chw(float const* dpuout_cls,
                           float const* dpuout_reg,
                           float const* class_threshold,
                           int const& max_result_size,
                           int filled_size_,
                           int const* type_map,
                           int const& net_length,
                           int const& net_width,
                           int const& preds_per_voxel,
                           int const& cls_number,
                           int const& reg_number,
                           int const& lane_feature_size,
                           float const* voxels_size,
                           float const* pts_range,
                           float* dense_pts_ptr)
{
    float class_threshold_sig[cls_number]{0};
    for (std::uint16_t i{0U}; i < cls_number; ++i)
    {
        class_threshold_sig[i] = -1 * log((1 - class_threshold[i]) / class_threshold[i]);
    }

    bool* visited = new bool[net_length * net_width * preds_per_voxel * cls_number];
    std::memset(visited, false, sizeof(bool) * net_length * net_width * preds_per_voxel * cls_number);

    float score{0};
    std::uint16_t label;
    int id{0};
    for (std::uint16_t i{0U}; i < net_length; ++i)
    {
        for (std::uint16_t j{0U}; j < net_width; ++j)
        {

            int index_{i * net_width + j};
            for (std::uint16_t k{0U}; k < preds_per_voxel; ++k)
            {
                for (std::uint16_t l{0U}; l < cls_number; ++l)
                {
                    // step1: decode score and label
                    std::uint64_t cur_index = static_cast<std::uint64_t>(index_ + (k * cls_number + l) * net_length * net_width);
                    score = dpuout_cls[cur_index];
                    label = l;

                    if(visited[cur_index])
                    {
                        continue;
                    }

                    if (score < class_threshold_sig[label])
                    {
                      continue;
                    }

                    // step2: decode points
                    if (filled_size_ >= max_result_size)
                    {
                        return filled_size_;
                    }
                    filled_size_ = dfs(i, j, k, l, id, dpuout_cls, dpuout_reg, class_threshold_sig, max_result_size, filled_size_, type_map, net_length,
                                       net_width, preds_per_voxel, cls_number, reg_number, lane_feature_size, voxels_size, pts_range, visited, dense_pts_ptr);
                    id++;

//                    int l_ = type_map[l];
//                    int cur_direction[2]{0};
//                    float cur_xy[2]{0.0};
//
//
//
//                    const float& cur_cw_distance = dpuout_reg[i * net_width + j + (k * reg_number * 2 + l_ * 2 + 0) * net_length * net_width];
//                    distance_decode(cur_cw_distance, true, cur_direction, cur_xy);
//                    std::cout << "cur_cw_distance: " << cur_cw_distance << std::endl;
//                    std::cout << "cur_xy[0]: " << cur_xy[0] << std::endl;
//                    std::cout << "cur_xy[1]: " << cur_xy[1] << std::endl;
//
//
//
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 0] = 0;
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 1] = (j + cur_xy[1]) * voxels_size[0] + pts_range[0];
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 2] = (i + cur_xy[0]) * voxels_size[1] + pts_range[1];
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 3] = 0;
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 4] = static_cast<float>(label);
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 5] = sigmoid(score);
//                    filled_size_++;
//
//                    const float& cur_ccw_distance = dpuout_reg[i * net_width + j + (k * reg_number * 2 + l_ * 2 + 1) * net_length * net_width];
//                    distance_decode(cur_ccw_distance, false, cur_direction, cur_xy);
//                    std::cout << "cur_ccw_distance: " << cur_ccw_distance << std::endl;
//                    std::cout << "cur_xy[0]: " << cur_xy[0] << std::endl;
//                    std::cout << "cur_xy[1]: " << cur_xy[1] << std::endl;
//
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 0] = 0;
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 1] = (j + cur_xy[1]) * voxels_size[0] + pts_range[0];
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 2] = (i + cur_xy[0]) * voxels_size[1] + pts_range[1];
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 3] = 0.5;
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 4] = static_cast<float>(label);
//                    dense_pts_ptr[filled_size_ * lane_feature_size + 5] = sigmoid(score);
//                    filled_size_++;
                }
            }
        }
    }
    delete[] visited;

    std::cout << "filled_size_: " << filled_size_ << std::endl;
    return filled_size_;
}

int decode_lane_head_double_cpu(py::array_t<float> out_cls,
                             py::array_t<float> out_reg,
                             py::array_t<int> type_map,
                             py::array_t<float> class_threshold,
                             py::array_t<float> voxels_size,
                             py::array_t<float> pts_ranges,
                             int max_result_size,
                             int filled_size_,
                             int preds_per_voxel,
                             int reg_number,
                             int lane_feature_size,
                             py::array_t<float> lane_dense_pts,
                             bool is_hwc)
 {
    // Get  ptr from python
    py::buffer_info cls_buf = out_cls.request();
    py::buffer_info reg_buf = out_reg.request();
    py::buffer_info type_map_buf = type_map.request();
    py::buffer_info cla_buf = class_threshold.request();
    py::buffer_info vox_buf = voxels_size.request();
    py::buffer_info pts_buf = pts_ranges.request();
    py::buffer_info dense_pts_buf = lane_dense_pts.request();

    int net_length = cls_buf.shape[1];
    int net_width  = cls_buf.shape[2];
    int cls_number = type_map_buf.shape[0];

    float* cls_ptr = (float*)cls_buf.ptr;
    float* reg_ptr = (float*)reg_buf.ptr;
    int* type_map_ptr = (int*)type_map_buf.ptr;
    float* cla_ptr = (float*)cla_buf.ptr;
    float* vox_ptr = (float*)vox_buf.ptr;
    float* pts_ptr = (float*)pts_buf.ptr;
    float* dense_pts_ptr = (float*)dense_pts_buf.ptr;

//    std::cout << "max_result_size: " << max_result_size << std::endl;
//    std::cout << "filled_size_: " << filled_size_ << std::endl;
//    std::cout << "preds_per_voxel: " << preds_per_voxel << std::endl;
//    std::cout << "reg_number: " << reg_number << std::endl;
//    std::cout << "lane_feature_size: " << lane_feature_size << std::endl;
    if (is_hwc)
    {
        std::cout << "Do not support now." << std::endl;
        return -1;
    }
    else
    {

        filled_size_ = decode_lane_double_chw(cls_ptr, reg_ptr, cla_ptr, max_result_size, filled_size_, type_map_ptr, net_length,
                                                net_width, preds_per_voxel, cls_number, reg_number, lane_feature_size, vox_ptr, pts_ptr,
                                                dense_pts_ptr);
    }


    return filled_size_;
 }
